package com.example.user.config;

import com.example.commons.util.ThrowableUtil;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.MDC;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.util.concurrent.ListenableFuture;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.Callable;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Future;

/**
 * Author:汪彬彬
 * date:2024/11/20 9:24
 */
@Slf4j
public class CustomThreadPoolExecutor extends ThreadPoolTaskExecutor {



        @Override
        public void execute(Runnable task) {
            super.execute(wrap(task));
        }

        @Override
        public void execute(Runnable task, long startTimeout) {
            super.execute(wrap(task), startTimeout);
        }

        @Override
        public Future<?> submit(Runnable task) {
            return super.submit(wrap(task));
        }

        @Override
        public <T> Future<T> submit(Callable<T> task) {
            return super.submit(wrap(task));
        }

        @Override
        public ListenableFuture<?> submitListenable(Runnable task) {
            return super.submitListenable(wrap(task));
        }

        @Override
        public <T> ListenableFuture<T> submitListenable(Callable<T> task) {
            return super.submitListenable(wrap(task));
        }

        @Override
        protected void cancelRemainingTask(Runnable task) {
            super.cancelRemainingTask(task);
        }

        private Runnable wrap(Runnable runnable) {
            Map<String, String> copyOfContextMap = MDC.getCopyOfContextMap();
            return () -> {
                MDC.setContextMap(copyOfContextMap);
                try {
                    runnable.run();
                } catch (Exception e) {
                    log.error("have error: " + ThrowableUtil.getStackTrace(e));
                } finally {
                    MDC.clear();
                }
            };
        }

        private <T> Callable<T> wrap(Callable<T> callable) {
            Map<String, String> copyOfContextMap = MDC.getCopyOfContextMap();
            return () -> {
                MDC.setContextMap(copyOfContextMap);
                try {
                    return callable.call();
                } catch (Exception e) {
                    log.error("have error: " + ThrowableUtil.getStackTrace(e));
                    // 后续的 future.get() 我们仍然期望能获取到异常，所以这里将异常抛出，并记录下这个异常，方便traceId跟踪日志
                    throw e;
                } finally {
                    MDC.clear();
                }
            };
        }


    }

